import matplotlib.pyplot as plt
from keras.datasets import fashion_mnist
import tensorflow as tf

tf.random.set_seed(18)

# 加载Fashion-MNIST数据集
(trax, tray), (tesx, tesy) = fashion_mnist.load_data()

# 将数据reshape为适合模型输入的形状并标准化到[0, 1]范围
trax = trax.reshape(-1, 28, 28, 1) / 255.0
tesx = tesx.reshape(-1, 28, 28, 1) / 255.0
print(trax.shape)

# 定义VGG16模型
from keras import Model, layers, Sequential, optimizers, activations, losses


class VGG16(Model):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.conv = Sequential([
            layers.Conv2D(filters=16, kernel_size=(3, 3), activation=activations.relu, padding='same'),
            layers.Conv2D(filters=16, kernel_size=(3, 3), activation=activations.relu, padding='same'),
            layers.MaxPooling2D(),
            layers.Conv2D(filters=32, kernel_size=(3, 3), activation=activations.relu, padding='same'),
            layers.Conv2D(filters=32, kernel_size=(3, 3), activation=activations.relu, padding='same'),
            layers.MaxPooling2D(),
            layers.Conv2D(filters=64, kernel_size=(3, 3), activation=activations.relu, padding='same'),
            layers.Conv2D(filters=64, kernel_size=(3, 3), activation=activations.relu, padding='same'),
            layers.Conv2D(filters=64, kernel_size=(3, 3), activation=activations.relu, padding='same'),
            layers.MaxPooling2D(),
            layers.Conv2D(filters=128, kernel_size=(3, 3), activation=activations.relu, padding='same'),
            layers.Conv2D(filters=128, kernel_size=(3, 3), activation=activations.relu, padding='same'),
            layers.Conv2D(filters=128, kernel_size=(3, 3), activation=activations.relu, padding='same'),
            layers.MaxPooling2D(),
        ])
        self.flat = layers.Flatten()
        self.func = Sequential([
            layers.Dense(units=64, activation=activations.relu),
            layers.Dropout(0.4),
            layers.Dense(units=64, activation=activations.relu),
            layers.Dropout(0.4),
            layers.Dense(units=10, activation=activations.softmax)  # Fashion-MNIST有10个类别
        ])

    def call(self, inputs, training=None, mask=None):
        out = self.conv(inputs)
        out = self.flat(out)
        out = self.func(out)
        return out


# 创建VGG16模型的实例
model = VGG16()

# 指定模型的输入形状
model.build(input_shape=[None, 28, 28, 1])

# 打印模型的概况信息
model.summary()

# 编译模型，选取合适的优化器、损失函数及准确率函数
choose = 2000
model.compile(optimizer=optimizers.Adam(), loss=losses.sparse_categorical_crossentropy, metrics='acc')
history = model.fit(trax[:choose], tray[:choose], epochs=10)

# 画出不同优化器的损失值
plt.plot(history.history['acc'], c='r', label='Adam')

# 使用RMSprop优化器
model = VGG16()
model.compile(optimizer=optimizers.RMSprop(), loss=losses.sparse_categorical_crossentropy, metrics='acc')
history = model.fit(trax[:choose], tray[:choose], epochs=10)
plt.plot(history.history['acc'], c='b', label='RMSprop')

# 使用Adagrad优化器
model = VGG16()
model.compile(optimizer=optimizers.Adagrad(), loss=losses.sparse_categorical_crossentropy, metrics='acc')
history = model.fit(trax[:choose], tray[:choose], epochs=10)
plt.plot(history.history['acc'], c='g', label='Adagrad')

plt.legend()
plt.show()
